# coding=utf-8


import os
import time, sys
import tensorflow as tf
from PIL import Image
import matplotlib.pyplot as plt
import threading

def get_target_classes(dir):
    out_class_dirs = []
    for class_dir in os.listdir(dir):
        if os.path.isdir(dir + "/" + class_dir):
            out_class_dirs.append(class_dir)
    return out_class_dirs


def get_finished_tf(dir):
    tf_list = []
    for tf_file in os.listdir(dir):
        if os.path.isfile(dir + "/" + tf_file) and os.path.splitext(tf_file)[1] == ".tfrecords":
            tf_list.append(tf_file)
    return tf_list


def get_real_target_classes(in_images_dirs, tf_files):
    images_dirs = []
    for im_dir in in_images_dirs:
        # print("checking ", im_dir)
        is_found = False
        for tf_file in tf_files:
            if str(im_dir) in str(tf_file):
                is_found = True
                break
        if is_found is False:
            images_dirs.append(im_dir)
    return images_dirs


def tf_write_one(writer, img_name, class_path,name):
    # print('路径',class_path+"/"+img_name)

    img_path = class_path + img_name  # 每一个图片的地址
    img = Image.open(img_path, 'r')
    size = img.size
    # print(size[1],size[0])
    # print(size)
    # print(img.mode)
    img_raw = img.tobytes()
    example = tf.train.Example(
        features=tf.train.Features(feature={
            'label': tf.train.Feature(bytes_list=tf.train.BytesList(value=[name.encode(encoding="utf-8")])),
            'img_raw': tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw])),
            'img_width': tf.train.Feature(int64_list=tf.train.Int64List(value=[size[0]])),
            'img_height': tf.train.Feature(int64_list=tf.train.Int64List(value=[size[1]]))
        }))
    writer.write(example.SerializeToString())


def tf_write(inpath, outpath):
    if os.path.exists(outpath) is False:
        os.makedirs(outpath)

    classes = get_target_classes(inpath)
    print("target classes:", classes)
    finished_tfs = get_finished_tf(outpath)
    print("finished classes:", finished_tfs)
    classes = get_real_target_classes(classes, finished_tfs)
    print("real target classes:", classes)

    version = str(int(time.time()))

    for index, name in enumerate(classes):
        print(index)
        print("生成：", name)
        class_path = inpath + '/' + name + '/'
        ftrecordfilename = ("%s.tfrecords-%s.tfrecords" % (version, name))
        writer = tf.python_io.TFRecordWriter(outpath + "/" + ftrecordfilename)
        for img_name in os.listdir(class_path):
            # tf_write_one(writer,img_name,class_path,name)
            t = threading.Thread(target=tf_write_one, args=(writer,img_name,class_path,name,))
            t.setDaemon(True)
            t.start()
        writer.close()


if __name__ == "__main__":

    if len(sys.argv) < 5:
        print("usage: ./classifier_tf_gen.py -i image_path -o train_vail_set")
        exit(-1)

    start_time = time.time()
    inpath = sys.argv[2]
    outpath = sys.argv[4]

    tf_write(inpath, outpath)
    print(' ALL COST: %d second' % (time.time() - start_time))
